#include "evdeploy/runtime/trt/trt_base.h"
#if defined(_M_ARM64) || defined(__aarch64__)
#define ONNX_PARSER_LIB ("/usr/lib/aarch64-linux-gnu/libnvonnxparser.so")
#define NV_INFER_LIB ("/usr/lib/aarch64-linux-gnu/libnvinfer.so")
#else
#define ONNX_PARSER_LIB ("/usr/lib/x86_64-linux-gnu/libnvonnxparser.so")
#define NV_INFER_LIB ("/usr/lib/x86_64-linux-gnu/libnvinfer.so")
#endif
namespace ev
{
    typedef void *(INFER_BUILDER_CREATOR1)(void *logger, int32_t version);

    EVTrtModel::EVTrtModel()
    {
        m_ev_engine = new EVCudaEngine;
        EVLOG(INFO) << "this is a tensorrt engine";
    }

    EVTrtModel::~EVTrtModel()
    {
        if (m_ev_context)
        {
            m_ev_context->destroy();
            delete m_ev_context;
            m_ev_context = nullptr;
        }
        if (m_ev_engine)
        {
            m_ev_engine->destroy();
            delete m_ev_engine;
            m_ev_engine = nullptr;
        }
    }

    void EVTrtModel::LoadTrt(const std::string &str_trt_name)
    {
        int ret = m_ev_engine->loadModel(str_trt_name);
        if (ret != 0)
        {
            EVLOG(ERROR) << "loading trt model wrong";
            exit(-1);
        }
        m_ev_context = m_ev_engine->createExecutionContext();

        if (m_ev_context == nullptr)
        {
            EVLOG(ERROR) << "null";
            exit(-1);
        }
        cudaStreamCreate(&m_cuda_stream);
        EVLOG(INFO) << "loading trt model done";
    }

    void EVTrtModel::ConvertEngine(const std::string str_onnx_name)
    {
        Logger gLogger;
        // tensorrt pipeline
        EVLOG(INFO) << "converting onnx model with custom mode";
        IBuilder *builder = nvinfer1::createInferBuilder(gLogger);
        builder->setMaxBatchSize(1);
        IBuilderConfig *config = builder->createBuilderConfig();
        config->setMaxWorkspaceSize(1ULL << 31);
        if (builder->platformHasFastFp16())
        {
            config->setFlag(BuilderFlag::kFP16);
            LOG(INFO) << "using FP16";
        }
        int ret = m_ev_engine->convertEngine(str_onnx_name, builder, config);
        if (ret != 0)
        {
            EVLOG(ERROR) << "converting trt model wrong";
            exit(-1);
        }
        config->destroy();
        builder->destroy();
    }

} // namespace ev
